import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from utils.features import prepare_for_training
# 原始数据
data_source_path = '../data/world-happiness-report-2017.csv'
source_data: pd.DataFrame = pd.read_csv(data_source_path)
# 训练数据和测试数据比例参数
frac: float = 0.8
# 训练数据 测试数据
features_cols = ['Economy..GDP.per.Capita.']
labels_cols = ['Happiness.Score']
# 预处理数据参数
polynomial_degree=0 #
sinusoid_degree=0
is_normalize=True

def get_xy_train(source_data,features_cols,labels_cols,frac=0.8):
    train_src_data: pd.DataFrame = source_data.sample(frac=train_frac)
    x_train: np.ndarray = train_src_data[features_cols].values
    y_train: np.ndarray = train_src_data[labels_cols].values
    # 获取测试数据
    test_src_data: pd.DataFrame = source_data.drop(train_src_data.index)
    x_test: np.ndarray = test_src_data[features_cols].values
    y_test: np.ndarray = test_src_data[labels_cols].values
    return x_train,y_train,x_test,y_test
# 预测数据
print("sdfsf")
def draw_plot(x_train,y_train,x_test, y_test,feature_column_names:str,label_column_names:str,
              label_train, label_test,title):
    plt.scatter(x_train, y_train, label=label_train)
    plt.scatter(x_test, y_test, label=label_test)
    plt.xlabel(feature_column_names)
    plt.ylabel(label_column_names)
    plt.title('Happy')
    plt.legend()
    plt.show()
def start():
    draw_plot(x_train,y_train,x_test, y_test,"".join(features_cols),"".join(labels_cols),
              label_train='Train data', label_test='Test data',title="Happy")

    (x_train,y_train,x_test,y_test)=get_xy_train(source_data,features_cols,labels_cols,frac)
    x_train:np.ndarray=x_train
    y_train:np.ndarray=y_train
    x_test:np.ndarray=x_test
    y_test:np.ndarray=y_test
    
    plt.close()

if __name__ == '__main__':
    start()
